package com.pudding.embeddingdemo;

import com.pudding.tools.ModelUtil;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.redis.RedisEmbeddingStore;

import java.util.List;

public class VectorDemo {
    public static void main(String[] args) {
        // 创建向量模型
        EmbeddingModel embeddingModel = OpenAiEmbeddingModel.builder()
                .baseUrl(ModelUtil.BASE_URL_OPENAI)
                .apiKey(ModelUtil.API_KEY_OPENAI)
                .build();

        // ⽣成向量
        Response<Embedding> embed = embeddingModel.embed("你好，我是楼兰");
        System.out.println(embed.content().toString());
        System.out.println(embed.content().vector().length);

        RedisEmbeddingStore embeddingStore = RedisEmbeddingStore.builder()
                .host("127.0.0.1")
                .port(6379)
                .dimension(1536) //维度，需要与计算结果保持⼀致。如果使⽤其他的模型，可能会有不同的结果。
                .build();

        // 存储向量--单独指定数据ID，索引的ID是固定的embedding-index
        embeddingStore.add("vec1", embed.content());

        Response<Embedding> embed1 = embeddingModel.embed("我的名字叫楼兰");
        // 去redis中找到最相似的，最多查3条，相似度下限为-1
        List<EmbeddingMatch<TextSegment>> matchedResult = embeddingStore.findRelevant(embed1.content(), 3, -1);
        for (EmbeddingMatch<TextSegment> embeddingMatch : matchedResult) {
            System.out.println(embeddingMatch.score());//打印相似度
        }
    }
}
